# code_sample/predo-merge-pid-timerdatas.py

#!/usr/bin/env python3
# -*- coding: utf-8 -*

import os
import re
import json
import argparse
import subprocess
from multiprocessing import Pool
import sys 

sys.path.append("..")
from package.utility import ensure_exists
from package.utility import find_targets
from package.utility import scan_runcases
from package.utility import load_runcase_npt_config
from package.utility import load_pidtimer_data
from package.utility import get_minmaxmean

'''
功能: 将每个算例目录下所有进程上的计时器合并到一个 json 文件中
输出格式:
{
    "timer_name_XXX1": {
    "level": level,
    "tmin": tmin,
    "tmax": tmax,
    "tavg": tavg
    },
    "timer_name_XXX2": {}
},
实现思路:
1. 找到所有的运行算例所在目录
2. 每个算例的合并只能由一个进程来完成, 因此这里的多进程任务分配以算例为单位, 而不是以 pidlog 为单位。
3. merge_pid_timerdatas 是具体的合并函数

日期: 2023-01-13T15:56
作者: 田鸿运
'''

def make_parser(): 
    NCPU=os.cpu_count()
    parser = argparse.ArgumentParser(
        prog="jxpamg-predo-merge-pid-timerdata",
        description="jxpamg timer data preprocess script, merge pid's timer data to a single json.")
    parser.add_argument(
            "LOGDIR",
            metavar="LOGDIR",
            help="Destination log directories."
        )
    parser.add_argument(
            "-p", "--proc",
            dest="num_proc",
            default=int(NCPU/2),
            help="Number of processes used to run the task, default:%(default)s"
        )
    return parser

def merge_pid_timerdatas(case_rundir):
    case_rundir = os.path.abspath(case_rundir)
    _, proc_num, _ =load_runcase_npt_config(case_rundir)
    pid_timer_logs = find_targets(case_rundir, "pid*.json")
    logs_num = len(pid_timer_logs)
    if logs_num != proc_num:
        print(f"WARNING: pid timer log numbers mismatch: found {logs_num}, exoected P{proc_num}")
        return
    
    output = case_rundir + "timers.json"
    if os.path.exists(output):
        os.remove(output)
    merged_timer = {}
    print(f">> {case_rundir}")
    for log in pid_timer_logs:
        data = load_pidtimer_data(log)
        for key,val in data.items():
            level_id, time_val = val
            if key not in merged_timer:
                merged_timer[key] = [level_id, [time_val]]
            else:
                merged_timer[key][1].append(time_val)

    merged_timer2 = {}
    for key, val in merged_timer.items():
        level_id, pid_timers = val
        tmin, tmax, tmean = get_minmaxmean(pid_timers)
        merged_timer2[key] = {
            "level": level_id,
            "tmix": tmin,
            "tmax": tmax,
            "tavg": tmean
        }
    print(f">> dump to {output}")
    with open(output, "w", encoding="utf-8") as fp:
        json.dump(merged_timer2, fp, indent=2, sort_keys=True)

def main():
    parser = make_parser()
    args = parser.parse_args()
    pool = Pool(int(args.num_proc))
    runcases = scan_runcases(args.LOGDIR)
    pool.map(merge_pid_timerdatas, runcases)
    print(">> DONE")

if __name__ == "__main__":
    main()
